from pulp import *
import time
import random
import matplotlib.pyplot as plt

# 生成多个size数据用于测量不同时间
def createData(sizes=[]):
    datasets = []
    for size in sizes:
        data = single_size_generate(size, size)
        datasets.append(data)
    return datasets

# 单次size生成数据
def single_size_generate(num_points, num_sets):
    # 生成X集合
    universe = set(range(num_points))
    # 生成解集合
    feasible_sets = []
    current_set = set(random.sample(universe, 20))
    feasible_sets.append(current_set)
    curNum = 1
    while len(universe - set().union(*feasible_sets)) >= 20 and curNum+1<num_sets:
        n = random.randint(1, min(20, num_points - len(current_set), len(current_set)))
        x = random.randint(1, n)
        next_set = set(random.sample(universe - current_set, x))
        next_set |= set(random.sample(current_set, n - x))
        current_set = next_set
        feasible_sets.append(current_set)
    feasible_sets.append(universe-set().union(*feasible_sets))
    # 生成其余集合
    remaining_sets = []
    while len(remaining_sets) < num_sets - len(feasible_sets):
        s = set(random.sample(universe, random.randint(1, num_points // 2)))
        if s not in feasible_sets and s not in remaining_sets:
            remaining_sets.append(s)

    # 合并可行解集合和其余集合
    subsets = feasible_sets + remaining_sets
    return universe, subsets

"""------------------------------------------线性规划算法---------------------------------------------------""" 
def lp_method(X, F):
    # 定义线性规划问题
    prob = LpProblem("SetCover", LpMinimize)
    # 定义变量，每个子集是否被选中
    set_vars = LpVariable.dicts("Set", range(len(F)), 0, 1, LpInteger)
    # 定义目标函数，最小化选中的子集数量
    prob += lpSum(set_vars[i] for i in range(len(F)))
    # 定义约束条件，每个元素必须被覆盖
    for x in X:
        prob += lpSum(set_vars[i] for i in range(len(F)) if x in F[i]) >= 1
    # 解决线性规划问题
    prob.solve()
    # 提取结果
    selected_sets = [F[i] for i in range(len(F)) if value(set_vars[i]) == 1]
    return selected_sets

"""------------------------------------------贪心算法---------------------------------------------------""" 
def greedy_method(X, F):
    # 初始化未覆盖元素集合和已选择的子集集合
    uncovered = set(X)
    selected_sets = []
    while uncovered:
        # 在未覆盖元素中找到能覆盖最多未覆盖元素的子集
        best_set = None
        max_covered = set()
        for subset in F:
            covered = subset.intersection(uncovered)
            if len(covered) > len(max_covered):
                max_covered = covered
                best_set = subset
        # 将找到的最优子集加入已选择的集合中，并更新未覆盖元素集合
        selected_sets.append(best_set)
        uncovered -= best_set
    return selected_sets


# 运行算法并测试时间
def run_algorithm(algorithm, data): 
    X,F = data
    start_time = time.time()
    algorithm(X,F)
    end_time = time.time()
    return (end_time - start_time) * 1000

# 画出不同size的运行时间变化图
def draw(greedy,LP, data_sets):
    # 注意算法运行顺序，此处写死
    # greedy贪心算法
    sizes1 = [len(data[0]) for data in data_sets]
    times1 = [run_algorithm(greedy, data) for data in data_sets]
    # LP线性规划算法
    sizes2 = [len(data[0]) for data in data_sets]
    times2 = [run_algorithm(LP, data) for data in data_sets]
    # 把两条算法曲线画在同个图
    plt.plot(sizes1, times1, '-o', label ='greedy')    # '-o'表示【线+大圆点】，也可以用'-.'表示【线+小圆点】
    plt.plot(sizes2, times2, '-o', label ='LP')
    plt.legend()
    plt.xlabel("Data Size")
    plt.ylabel("Time (ms)")
    plt.title("comparation of Greedy and LP method")
    #plt.savefig("comparation.png") # 用于保存图片
    plt.show()

# 测试用的
#X = set([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
#F = [set([1, 2, 3]), set([2, 3, 4]), set([4, 5, 6]), set([ 8, 9, 10]), set([7])]

# 生成数据
sizes= list([100,200,300,400,500,1000,2000])
data = createData(sizes)

draw(greedy_method,lp_method, data)
